package graph

import (
    "fmt";
)



type uset map[Vertex] bool;

var nilset uset = nil;

type UnionFind struct {
    components map[Vertex] uset ;
    p map[Vertex] Vertex;
    size map[Vertex]int;
}


func NewUnionFind(nodes VertexIterator) (*UnionFind) {
    components := make(map[Vertex] uset);
    p := make(map[Vertex] Vertex);
    size := make(map[Vertex] int);
    for !nodes.End() {
        nodes.Next();
        p[nodes.Current()] = nodes.Current();
        size[nodes.Current()] = 1;
    }
    return &UnionFind { components, p, size };
}

// This is combined operation. It is essentially Union(Find(u), Find(v))
func (uf *UnionFind) FindUnion(u Vertex, v Vertex) {
    if u == v {
        return;
    }
    r1 := uf.Find(u);
    r2 := uf.Find(v);
    uf.Union(r1, r2);
}

func (uf *UnionFind) Union(r1 Vertex, r2 Vertex) {
    size1 := uf.size[r1];
    size2 := uf.size[r2];
    if (size1 > size2) {
        uf.size[r1] = size1 + size2;
        uf.p[r2] = r1;
       // uf.addToComponents(r1, r2);
    } else {
        uf.size[r2] = size1 + size2;
        uf.p[r1] = r2;
        //uf.addToComponents(r2, r1);
    }
}

func (uf *UnionFind) addToComponents(parent Vertex, child Vertex) {
    parentSet, present := uf.components[parent];
    if !present {
        parentSet = make(map[Vertex] bool);
        uf.components[parent] = parentSet;
    }
    
    childSet, present := uf.components[child];
    if present {
        for k,v := range childSet {
            parentSet[k] = v;
        }
        uf.components[child] = nil;
    }
    fmt.Printf("%s %s\n", parentSet, child);
    parentSet[child] = true;
}

func (uf *UnionFind) Find(n Vertex) (root Vertex){
    if uf.p[n] == n {
        return n;
    }
    return uf.Find(uf.p[n]);
}

func (uf *UnionFind) SameComponents(u Vertex, v Vertex) bool {
    return uf.Find(u) == uf.Find(v);
}

// Returns a array containing of all the nodes that are in 
// same component as node u.
func (uf *UnionFind) GetComponent(u Vertex) []Vertex {
    root := uf.Find(u);
    set, present := uf.components[root];
    if !present {
        return [] Vertex { u };
    }
    list := make([]Vertex, len(set) + 1, len(set) + 1);
    i := 0;
    for node,_ := range set {
        list[i] = node;
        i++;
    }
    list[i] = root;
    return list;
    
}


